#include "libmri.h"

void calib3d_orthiter (complex *smap, float *emap, complex *KKt,
                       int *xsize, int *ksize, int qcols, int qriters)
{
	int sx1 = xsize[0], sx2 = xsize[1], sx3 = xsize[2], nch = xsize[3];
	int ks1 = ksize[0], ks2 = ksize[1], ks3 = ksize[2];
	int kk1 = 2*ks1-1, kk2 = 2*ks2-1, kk3 = 2*ks3-1;

	complex *km = new complex [sx1*sx2*sx2*nch*nch];

	for (int i = 0; i < sx1*sx2*sx2*nch*nch; ++i){
		km[i].real = km[i].imag = 0.0f;
	}

	int o1 = pad_index (sx1, kk1),
	    o2 = pad_index (sx2, kk2),
	    o3 = pad_index (sx3, kk3);

	#pragma omp parallel for
	for (int c2 = 0; c2 < nch; ++c2){
		for (int c1 = 0; c1 < nch; ++c1)
			for (int k3 = 0; k3 < kk3; ++k3)
				for (int k2 = 0; k2 < kk2; ++k2)
					for (int k1 = 0; k1 < kk1; ++k1)
					{
						complex &KKt_ = KKt [kk1-1-k1 + kk1*
						                    (kk2-1-k2 + kk2*
						                    (kk3-1-k3 + kk3*(c1 + nch*c2)))];

						complex &km_ = km [o1+k1 + sx1*
						                  (o2+k2 + sx2*
						                  (o3+k3 + sx3*(c1 + nch*c2)))];

						km_ = complex_multiply (KKt_,
						          compose_complex (sqrt(sx1*sx2*sx3), 0));
					}
	}

	int kmsize[4] = {sx1, sx2, sx3, nch*nch};
	
	fftwf_plan_with_nthreads (8);
	ifft3c ((cplxf*) km, kmsize);

	acmlsetnumthreads (1);
	#pragma omp parallel num_threads(8)
	{
		complex *A = new complex [nch*nch],
		        *Q = new complex [nch*qcols],
		        *T = new complex [nch*qcols],
		        *v = new complex [qcols];

		complex _1 [1] = {{1.0f, 0.0f}}, _0 [1] = {{0.0f, 0.0f}},
		        _n1[1] = {{-1.0f, 0.0f}};

		#pragma omp for
		for (int iz = 0; iz < sx3; ++iz)
		{
		  for (int iy = 0; iy < sx2; ++iy){
			for (int ix = 0; ix < sx1; ++ix)
			{
				for (int c2 = 0; c2 < nch; ++c2)
					for (int c1 = 0; c1 < nch; ++c1)
						A [c1 + nch*c2] = km [ix + sx1*
						                     (iy + sx2*
						                     (iz + sx3*(c1 + nch*c2)))];

				for (int c2 = 0; c2 < qcols; ++c2)
					for (int c1 = 0; c1 < nch; ++c1)
						Q [c1 + nch*c2].real = Q [c1 + nch*c2].imag = 0;

				for (int c1 = 0; c1 < qcols; ++c1)
					Q [c1 + nch*c1].real = 1.0f;

				for (int iter = 1; iter <= qriters; ++iter)
				{
					cgemm ('N', 'N', nch, qcols, nch, _1, A, nch, Q, nch, _0, T, nch);

					for (int j = 0; j < qcols; ++j){
						complex nrm = {1.0f/scnrm2 (nch, &T [nch*j], 1), 0.0f};
						cscal (nch, &nrm, &T [nch*j], 1);
					}

					for (int j = 1; j < qcols; ++j){
						cgemv ('C', nch,j, _1,  T,nch, &T[j*nch],1, _0, v,1);
						cgemv ('N', nch,j, _n1, T,nch, v,1, _1, &T[j*nch],1);
					}

					complex *tmp = Q; Q = T; T = tmp;
				}

				cgemm ('N', 'N', nch, qcols, nch, _1, A, nch, Q, nch, _0, T, nch);

				for (int j = 0; j < qcols; ++j)
					emap [ix + sx1*(iy + sx2*(iz + sx3*j))]
							= scnrm2 (nch, &T [nch*j], 1);

				for (int c2 = 0; c2 < qcols; ++c2)
					for (int c1 = 0; c1 < nch; ++c1)
						smap [ix + sx1*
						     (iy + sx2*
						     (iz + sx3*(c1 + nch*c2)))] = Q [c1 + nch*c2];
			}
		  }
		}

		delete [] A;
		delete [] Q;
		delete [] T;
		delete [] v;
	}

	delete [] km;
}
